import requests
import json

class OllamaClient:
    def __init__(self, base_url="http://localhost:11434"):
        self.base_url = base_url
        self.api_chat = f"{base_url}/api/chat"  # 使用 /api/chat 端点
        # 初始化消息列表，包含系统提示
        self.messages = [
            {
                'role': 'system',
                'content': """你能根据给你的单词，补充缺失的单词，连成完整的句子，并给出最后的句子"""
            }
        ]
        
    def set_system_prompt(self, new_prompt):
        """设置新的system prompt"""
        self.messages[0]['content'] = new_prompt
        
    def generate_response(self, prompt, model="qw:latest", stream=False):
        """
        向Ollama发送请求并获取响应
        
        Args:
            prompt (str): 输入的提示文本
            model (str): 使用的模型名称
            stream (bool): 是否使用流式输出
        """
        try:
            # 添加用户消息
            self.messages.append({
                'role': 'user',
                'content': prompt
            })
            
            data = {
                "model": model,
                "messages": self.messages,  # 使用消息格式
                "stream": stream
            }
            
            response = requests.post(
                self.api_chat,  # 使用 /api/chat 端点
                json=data,
                stream=stream
            )
            
            # 检查HTTP状态码
            if response.status_code != 200:
                print(f"请求失败，状态码: {response.status_code}")
                return None
            
            if stream:
                # 流式处理响应
                full_response = ""
                for line in response.iter_lines():
                    if line:
                        decoded_line = line.decode('utf-8')
                        json_response = json.loads(decoded_line)
                        if 'message' in json_response and 'content' in json_response['message']:
                            chunk = json_response['message']['content']
                            full_response += chunk
                            print(chunk, end='', flush=True)
                return full_response
            else:
                # 一次性返回完整响应
                return response.json()['message']['content']
                
        except requests.exceptions.RequestException as e:
            print(f"连接错误: {e}")
            return None
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            return None
        except Exception as e:
            print(f"发生错误: {e}")
            return None

def main():
    # 使用示例
    client = OllamaClient()
    
    # 可以设置新的system prompt
    custom_prompt = """你能根据给你的单词，补充单词，连成完整的句子，只给出最后的结果。回答的结果20字以内。"""
    client.set_system_prompt(custom_prompt)
    
    # 测试对话
    prompts = [
        "我 名字",
        "我 电话",
        "厕所 那",
        "厕所 哪"
    ]
    
    for prompt in prompts:
        print(f"\n用户: {prompt}")
        print("AI助手: ", end='')
        response = client.generate_response(prompt, stream=True)
        print("\n" + "="*50)

if __name__ == "__main__":
    main()